#include <stdint.h>
#include <stdio.h>
#include <assert.h>
#include <string.h>
#include "bgc2.h"
#include "meta_io.h"
#include "io_util.h"
#include "../config_vars.h"
#include "../check_syscalls.h"
#include "../universal_constants.h"
#include "../rockstar.h"
#include "../groupies.h"

char **bgc2_snapnames = NULL;
int64_t num_bgc2_snaps = 0;
GROUP_DATA_RMPVMAX *gd = NULL;


void populate_header(struct bgc2_output_header *hdr, int64_t id_offset, 
		     int64_t snap, int64_t chunk, float *bounds) {
  memset(hdr, 0, sizeof(struct bgc2_output_header));
  hdr->magic = BGC_MAGIC;
  hdr->version = 2;
  hdr->num_files = (PARALLEL_IO) ? NUM_WRITERS : 1;
  hdr->file_id = chunk;
  hdr->snapshot = snap;
  hdr->format = PDATA_FORMAT_PV;
  hdr->format_gdata = GDATA_FORMAT_RMPVMAX;
  hdr->first_valid_group_id = 0;
  hdr->first_group_id = id_offset;
  hdr->ngroups = 0;
  hdr->ngroups_total = 0;
  hdr->min_group_part = MIN_HALO_OUTPUT_SIZE;

  hdr->npart = 0;
  hdr->npart_total = 0;
  hdr->npart_orig = num_p;
  hdr->valid_part_ids = 1;

  hdr->max_npart = 0;
  hdr->max_npart_total = 0;

  hdr->linkinglength = FOF_LINKING_LENGTH;
  hdr->time = SCALE_NOW;
  hdr->redshift = (SCALE_NOW>0) ? (1.0/(SCALE_NOW) - 1.0) : 1e10;
  hdr->Omega0 = Om;
  hdr->OmegaLambda = Ol;
  hdr->BoxSize = BOX_SIZE;

  if (bounds) for (int64_t i=0; i<6; i++) hdr->bounds[i] = bounds[i];

  hdr->Hubble0 = h0;
  hdr->GravConst = Gc;
  hdr->part_mass = PARTICLE_MASS;
}

void output_bgc2(int64_t id_offset, int64_t snap, int64_t chunk, float *bounds)
{
  char buffer[1024];
  int64_t i, j, k, id=0;
  FILE *output;
  struct bgc2_output_header hdr = {0};
  struct particle **sphere = NULL;
  int64_t num_results;
  PARTICLE_DATA_PV *pd = NULL;
  int64_t num_to_print = count_halos_to_print(bounds);

  if (bgc2_snapnames == NULL) {
    if (!strlen(BGC2_SNAPNAMES)) return;
    read_input_names(BGC2_SNAPNAMES, &bgc2_snapnames, &num_bgc2_snaps);
  }

  for (i=0; i<num_bgc2_snaps; i++)
    if ((snapnames && !strcmp(snapnames[snap], bgc2_snapnames[i])) ||
	(!snapnames && atoi(bgc2_snapnames[i])==snap)) break;
  if (i==num_bgc2_snaps) return;

  assert(BGC2_HEADER_SIZE == sizeof(hdr));
  get_output_filename(buffer, 1024, snap, chunk, "bgc2");
  output = check_fopen(buffer, "w");

  populate_header(&hdr, id_offset, snap, chunk, bounds);
  hdr.ngroups = num_to_print;
  fwrite_fortran(&hdr, BGC2_HEADER_SIZE, 1, output);
  if (!num_to_print) {
    fclose(output);
    return;
  }

  gd = check_realloc(gd, sizeof(GROUP_DATA_RMPVMAX)*num_to_print,
		     "Allocating output halo buffer");

  fwrite_fortran(gd, sizeof(GROUP_DATA_RMPVMAX), num_to_print, output);

  build_particle_tree();
  for (i=0; i<num_halos; i++) {
    if (!_should_print(halos+i, bounds)) continue;
    sphere = find_halo_sphere(halos+i, &num_results);
    gd[id].id = id+id_offset;
    gd[id].parent = -1;
    gd[id].num_part = num_results;
    gd[id].radius = halos[i].r;
    gd[id].mass = halos[i].m;
    for (j=0; j<3; j++) {
      gd[id].pos[j] = halos[i].pos[j];
      gd[id].vel[j] = halos[i].pos[j+3];
    }

    hdr.npart += num_results;
    if (num_results > hdr.max_npart) {
      pd = check_realloc(pd, sizeof(struct particle)*num_results,
			 "Allocating particle output buffer.");
      hdr.max_npart = num_results;
    }

    for (j=0; j<num_results; j++) {
      pd[j].part_id = sphere[j]->id;
      for (k=0; k<3; k++) {
	pd[j].pos[k] = sphere[j]->pos[k];
	pd[j].vel[k] = sphere[j]->pos[k+3];
      }
    }

    fwrite_fortran(pd, sizeof(PARTICLE_DATA_PV), num_results, output);
    id++;
  }
  clear_particle_tree();

  rewind(output);
  fwrite_fortran(&hdr, BGC2_HEADER_SIZE, 1, output);
  fwrite_fortran(gd, sizeof(GROUP_DATA_RMPVMAX), num_to_print, output);

  free(pd);
  gd = check_realloc(gd, 0, "Freeing group data.");
  fclose(output);
}


void load_bgc2_groups(char *filename, struct bgc2_output_header *hdr,
		    GROUP_DATA_RMPVMAX **groups, int64_t *num_groups)
{
  FILE *input;
  int64_t new_group_size;

  assert(sizeof(struct bgc2_output_header) == BGC2_HEADER_SIZE);
  input = check_fopen(filename, "rb");

  fread_fortran(hdr, BGC2_HEADER_SIZE, 1, input, 0);
  assert(hdr->magic == BGC_MAGIC);
  assert(hdr->version == 2);
  assert(hdr->format_gdata == GDATA_FORMAT_RMPVMAX);

  new_group_size = sizeof(GROUP_DATA_RMPVMAX)*((*num_groups)+hdr->ngroups);
  *groups = check_realloc(*groups, new_group_size, "Allocating groups.");
  fread_fortran((*groups) + (*num_groups), sizeof(GROUP_DATA_RMPVMAX), 
		hdr->ngroups, input, 0);
  *num_groups += hdr->ngroups;
  fclose(input);
}

#define FAST3TREE_TYPE GROUP_DATA_RMPVMAX
#define FAST3TREE_PREFIX BGC2
#include "../fast3tree.c"

#define GROUP_LIST gd
#define RADIUS radius
#include "../parents.c"

int sort_by_id(const void *a, const void *b) {
  const GROUP_DATA_RMPVMAX *c = a;
  const GROUP_DATA_RMPVMAX *d = b;
  if (c->id < d->id) return -1;
  if (c->id > d->id) return 1;
  return 0;
}

void calc_bgc2_parents(int64_t snap)
{
  int64_t i, num_groups = 0, groups_written = 0;
  struct bgc2_output_header *hdrs = NULL;
  char buffer[1024];
  FILE *output;

  hdrs = check_realloc(hdrs, BGC2_HEADER_SIZE*NUM_WRITERS,
		       "Allocating BGC2 headers.");
  for (i=0; i<NUM_WRITERS; i++) {
    get_output_filename(buffer, 1024, snap, i, "bgc2");
    load_bgc2_groups(buffer, hdrs + i, &gd, &num_groups);
    if (!i) hdrs[0].npart_total = hdrs[0].ngroups_total = 0;
    hdrs[0].npart_total += hdrs[i].npart;
    hdrs[0].ngroups_total += hdrs[i].ngroups;
    if (hdrs[i].max_npart > hdrs[0].max_npart_total)
      hdrs[0].max_npart_total = hdrs[i].max_npart;
  }

  for (i=1; i<NUM_WRITERS; i++) {
    hdrs[i].npart_total = hdrs[0].npart_total;
    hdrs[i].ngroups_total = hdrs[0].ngroups_total;
    hdrs[i].max_npart_total = hdrs[0].max_npart_total;
  }

  find_parents(num_groups);
  qsort(gd, num_groups, sizeof(GROUP_DATA_RMPVMAX), sort_by_id);

  for (i=0; i<NUM_WRITERS; i++) {
    assert(hdrs[i].first_group_id == gd[groups_written].id);
    get_output_filename(buffer, 1024, snap, i, "bgc2");
    output = fopen(buffer, "r+b");
    fwrite_fortran(hdrs + i, BGC2_HEADER_SIZE, 1, output);
    fwrite_fortran(gd + groups_written, sizeof(GROUP_DATA_RMPVMAX), 
		   hdrs[i].ngroups, output);
    groups_written += hdrs[i].ngroups;
    fclose(output);
  }

  gd = check_realloc(gd, 0, "Freeing group data.");
  free(hdrs);
}
